{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, read_json\n",
    "import os\n",
    "from suno_utils.audio import Audio\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = read_jsonl(\"/app2/suno/data/sara/sfx_get_beats/combined_v3_w_extreme_metas_v0.jsonl\")\n",
    "metas_map = {meta[\"id\"]: meta for meta in meta}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bpm_std_to_time_difference(tempo_bpm, std_bpm):\n",
    "    # Beat periods in seconds\n",
    "    mean_period = 60 / tempo_bpm\n",
    "    slower_period = 60 / (tempo_bpm - 2 * std_bpm)  # tempo - 1 std dev\n",
    "    faster_period = 60 / (tempo_bpm + 2 * std_bpm)  # tempo + 1 std dev\n",
    "    \n",
    "    # Time differences from mean\n",
    "    time_diff_slower = slower_period - mean_period\n",
    "    time_diff_faster = mean_period - faster_period\n",
    "    \n",
    "    return time_diff_slower, time_diff_faster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "kept = 0\n",
    "total_found = 0\n",
    "\n",
    "folder_path = \"/app2/suno/data/sara/sfx_get_beats/sfx_full_beat_data/\"\n",
    "\n",
    "for filename in os.listdir(folder_path):\n",
    "    if filename.endswith('.json'):\n",
    "        file_path = os.path.join(folder_path, filename)\n",
    "        data = read_json(file_path)\n",
    "        for id, row in data.items():\n",
    "            meta_row = metas_map[id]\n",
    "            meta_row[\"beat_start_time_s\"] = 0\n",
    "            if not row[\"processing_failed\"]:\n",
    "                detected_bpm = row[\"inferred_tempo\"]\n",
    "                beats_detected = row[\"beats_detected\"]\n",
    "                beat_times = row[\"beat_times\"]\n",
    "                tempo_std = row[\"tempo_std\"]\n",
    "\n",
    "                if detected_bpm is not None and beats_detected is not None and tempo_std is not None:\n",
    "                    std_time = max(bpm_std_to_time_difference(detected_bpm, tempo_std))\n",
    "                    total_found += 1\n",
    "                    if beats_detected >= 4 and std_time < 0.01:\n",
    "                        meta_row[\"beat_start_time_s\"] = float(beat_times[0])\n",
    "                        kept += 1\n",
    "\n",
    "print(total_found)\n",
    "print(kept / total_found)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "valid = []\n",
    "\n",
    "for id, data in metas_map.items():\n",
    "    if \"beat_start_time_s\" in data and data[\"beat_start_time_s\"] > 0:\n",
    "        valid.append(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for v in tqdm(valid):\n",
    "    if v[\"dataset\"] != \"extreme\":\n",
    "        s3_path = v[\"s3_filepath\"]\n",
    "        beat_start_time_s = v['beat_start_time_s']\n",
    "        audio = Audio.from_s3(s3_path)\n",
    "        trim_audio = audio.get_slice(from_s=beat_start_time_s)\n",
    "        token_len_samples = int(trim_audio.sample_rate * 0.04)\n",
    "        audio_len_samples = int(trim_audio.sample_rate * trim_audio.duration_s)\n",
    "        leftover_samples = audio_len_samples % token_len_samples\n",
    "        target_len_samples = audio_len_samples - leftover_samples\n",
    "        if leftover_samples > 0:\n",
    "            target_len_samples += token_len_samples\n",
    "\n",
    "        target_len_seconds = target_len_samples / trim_audio.sample_rate\n",
    "        if target_len_samples > audio_len_samples:\n",
    "            trim_audio = trim_audio.pad_to_length(target_len_seconds)\n",
    "\n",
    "        if abs(trim_audio.duration_s - audio.duration_s) > 0.005 or s3_path.endswith(\".mp3\"):\n",
    "            output_path = os.path.join(\"/app2/suno/data/sara/sfx_get_beats/aligned_audio\", v[\"id\"] + \".opus\")\n",
    "            trim_audio.write_opus(output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k,v in valid[0].items():\n",
    "    print(k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}